#pragma once
#pragma once
#pragma once

#define RET_OK nullptr
#define USE_CUDA

#ifdef _WIN32
#include <Windows.h>
#include <direct.h>
#include <io.h>
#endif

#include <string>
#include <vector>
#include <cstdio>
#include <opencv2/opencv.hpp>
#include "onnxruntime_cxx_api.h"

#ifdef USE_CUDA
#include <cuda_fp16.h>
#endif

// 模型任务模式
enum MODEL_TYPE
{
    //FLOAT32 MODEL
    DETECT = 1,
    POSE = 2,
    CLASS = 3,

    //FLOAT16 MODEL
    DETECT_HALF = 4,
    POSE_HALF = 5,
    CLASS_HALF = 6
};

// 初始化参数
typedef struct _DL_INIT_PARAM
{
    std::string modelPath; // 模型路径
    std::vector<int> imgSize = { 640, 640 }; // 输入图像大小
    MODEL_TYPE modelType = DETECT; // 模式
    float rectConfidenceThreshold = 0.6; // 矩形框置信度
    float iouThreshold = 0.5;
    int	keyPointsNum = 2; // 关键点数量
    bool cudaEnable = true; // cuda启动标志位
    int logSeverityLevel = 3; // 推理日志的报错记录等级
    int intraOpNumThreads = 2; // 推理线程数
} DL_INIT_PARAM;

// 输出接口(推理结果)
typedef struct _DL_RESULT
{
    int classId; // 类别
    float confidence; // 置信度
    cv::Rect box; // 矩形框
    std::vector<cv::Point2f> keyPoints; // 关键点
} DL_RESULT;

// 模型及推理
class myModel
{
public:
    myModel();

    ~myModel();

public:

    /*只需要使用CreateSession与RunSession两个函数*/
    const char* CreateSession(const DL_INIT_PARAM& iParams); // 创建推理环境与会话

    char* RunSession(cv::Mat& iImg, std::vector<DL_RESULT>& oResult); // 推理接口

    char* WarmUpSession(); // Cuda预热

    template<typename N>
    char* TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std::vector<int64_t>& inputNodeDims,
        std::vector<DL_RESULT>& oResult); // 推理与后处理

    char* PreProcess(cv::Mat& iImg, std::vector<int> iImgSize, cv::Mat& oImg); // 预处理

    std::vector<std::string> classes{}; // 任务类别
    DL_INIT_PARAM params; // 初始参数
    std::vector<DL_RESULT> res; // 输出结果

private:
    Ort::Env env; // 推理环境
    Ort::Session* session; // 推理会话
    bool cudaEnable;
    Ort::RunOptions options;
    std::vector<const char*> inputNodeNames; // 模型输入名
    std::vector<const char*> outputNodeNames; // 模型输出名

    MODEL_TYPE modelType; // 模型任务模式
    std::vector<int> imgSize; // 图片大小
    float rectConfidenceThreshold; // 矩形框置信度阈值
    float iouThreshold;
    float resizeScales; // letterbox scales
};
